# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import csv
import sys
import os
import pathlib
from collections import namedtuple

from absl import (flags, app)
import jinja2
import pandas

from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.exchange.okex_futures.kr_rest.currency import OkexCurrency
from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.dumper_base import run_from_archive
from coin.support.feed_tool.feed_stats.logic.util import *

FLAGS = flags.FLAGS

CsvWriter = namedtuple('CsvWriter', ['writer', 'file'])


class DumpStrategy(object):
  def __init__(self, products, csv_root, trading_date):
    self._products = products
    self._csv_root = os.path.expanduser(csv_root)
    pathlib.Path(csv_root).mkdir(parents=True, exist_ok=True)
    self._trading_date = trading_date
    self._csv_writers = {}

  def prepare_csv_file(self, product):
    csv_file = open(os.path.join(
        self._csv_root,
        'ticker_%s_%s.csv' % (product.symbol, self._trading_date.strftime('%Y%m%d'))),
                    'w',
                    encoding='utf-8')
    writer = csv.writer(csv_file, dialect='excel')
    csv_header = ['timestamp', 'vol_24h', 'open_interest', 'limit_low', 'limit_high']
    writer.writerow(csv_header)
    csv_file.flush()
    return writer, csv_file

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      if book_builder_name == get_subscription_key(product):
        book_builder.subscribe_ticker(product, self.on_ticker)
        exchange, symbol = to_book_builder_exchange_symbol(product)
        writer, csv_file = self.prepare_csv_file(product)
        self._csv_writers[(exchange, symbol)] = CsvWriter(writer, csv_file)

  def on_ticker(self, ticker):
    exchange, symbol = to_book_builder_exchange_symbol(ticker.product)
    writer = self._csv_writers.get((exchange, symbol), None)
    assert writer is not None

    timestamp = ticker.timestamp
    csv_line = [
        timestamp, ticker.vol_24h, ticker.open_interest, ticker.limit_low, ticker.limit_high
    ]
    writer.writer.writerow(csv_line)
    writer.file.flush()

  def close_csvs(self):
    for writer in self._csv_writers.values():
      writer.file.flush()
      writer.file.close()


class DbDumpStrategy(object):
  def __init__(self, products, csv_root, trading_date, query_machine):
    self._products = products
    self._csv_root = os.path.expanduser(csv_root)
    pathlib.Path(csv_root).mkdir(parents=True, exist_ok=True)
    self._trading_date = trading_date
    self._query_machine = query_machine

  def extract_feed_from_db_and_dump_to_csv(self):
    for product in self._products:
      csv_file = open(os.path.join(
          self._csv_root,
          'long_short_position_account_ratio_%s_%s_%s.csv' %
          (self._trading_date.strftime('%Y%m%d'), self._query_machine, product.symbol)),
                      'w',
                      encoding='utf-8')
      writer = csv.writer(csv_file, dialect='excel')
      csv_header = ['fetch_time', 'long_short_position_account_ratio']
      writer.writerow(csv_header)
      long_short_position_account_ratio = self.query_long_short_position_account_ratio()
      for _, row in long_short_position_account_ratio.iterrows():
        product_type = row['product_type']
        symbol = row['symbol']
        if symbol != product.symbol:
          continue
        fetch_time = row['fetch_time']
        data = eval(row['json'])['data']
        buy = data['buydata'][-1]
        sell = data['selldata'][-1]
        ratio = buy / sell
        writer.writerow([fetch_time, ratio])
      csv_file.flush()
      csv_file.close()

      csv_file = open(os.path.join(
          self._csv_root,
          'long_short_margin_ratio_%s_%s_%s.csv' %
          (self._trading_date.strftime('%Y%m%d'), self._query_machine, product.symbol)),
                      'w',
                      encoding='utf-8')
      writer = csv.writer(csv_file, dialect='excel')
      csv_header = ['fetch_time', 'long_short_margin_ratio']
      writer.writerow(csv_header)
      long_short_margin_ratio = self.query_long_short_margin_ratio()
      for _, row in long_short_margin_ratio.iterrows():
        product_type = row['product_type']
        symbol = row['symbol']
        if symbol != product.symbol:
          continue
        fetch_time = row['fetch_time']
        data = eval(row['json'])['data']
        buy = data['buydata'][-1]
        sell = data['selldata'][-1]
        ratio = buy / sell
        writer.writerow([fetch_time, ratio])
      csv_file.flush()
      csv_file.close()

  def _gen_query_column_sql(self, column):
    template_loader = jinja2.FileSystemLoader(searchpath="./")
    template_env = jinja2.Environment(loader=template_loader)

    sql_tmpl = 'coin/support/feed_tool/feed_stats/query_sql/' \
               'query_okex_instrument_feed.sql.tmpl'
    template = template_env.get_template(sql_tmpl)
    return template.render(trading_date=self._trading_date,
                           machine=self._query_machine,
                           column=column)

  def query_long_short_position_account_ratio(self):
    engine = gen_engine_for_instrument_feed_querier()
    feed = pandas.read_sql(self._gen_query_column_sql('long_short_position_account_ratio'), engine)
    return feed

  def query_long_short_margin_ratio(self):
    engine = gen_engine_for_instrument_feed_querier()
    feed = pandas.read_sql(self._gen_query_column_sql('long_short_margin_ratio'), engine)
    return feed


def launch(trading_date, ticker_feed_machine, query_machine, products, hours, csv_root):
  start_time = trading_date  # UTC 0 hour

  strategy = DumpStrategy(products, csv_root, trading_date)

  run_from_archive(products,
                   strategy.on_book_reset,
                   start_time,
                   start_time + datetime.timedelta(hours=hours),
                   machine=[ticker_feed_machine])

  strategy.close_csvs()

  currencys = [product.base for product in products]
  strategy = DbDumpStrategy(currencys, csv_root, trading_date, query_machine)
  strategy.extract_feed_from_db_and_dump_to_csv()

  return 0


def main(argv):
  trading_date = FLAGS.trading_date
  ticker_feed_machine = FLAGS.ticker_feed_machine
  query_machine = FLAGS.query_machine
  csv_root = FLAGS.csv_root
  assert trading_date, '--trading_date must be specified.'
  assert ticker_feed_machine, '--ticker_feed_machine must be specified.'
  assert query_machine, '--query_machine must be specified.'
  assert csv_root, '--csv_root must be specified.'

  symbols = flags.FLAGS.symbols.split(",")
  print('Running for %s %s ...' % (trading_date, ticker_feed_machine))
  sys.stdout.flush()
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')
  products = [OkexFuturesProduct.FromStr(symbol, trading_date) for symbol in symbols]
  launch(trading_date, ticker_feed_machine, query_machine, products, FLAGS.hours, csv_root)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('ticker_feed_machine', None, 'Instance machine name.')

  flags.DEFINE_string('query_machine', None, 'Instance machine name.')

  flags.DEFINE_integer('hours', 24, 'Time span in hours from 00:00:00 UTC.')

  flags.DEFINE_string('csv_root', None, 'Output csv files root directory.')

  flags.DEFINE_string('symbols', 'BTC-USD.QUARTER,ETH-USD.QUARTER,EOS-USD.QUARTER', 'symbol name')

  app.run(main)
